// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/serializers/unsafe_row_serializer.h>
#include <gtest/gtest.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/row/unsafe_row_fast.h>
#include <pollux/vector/vector_fuzzer.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo;
using namespace kumo::pollux;

struct TestParam {
    common::CompressionKind compressionKind;
    bool appendRow;

    TestParam(common::CompressionKind _compressionKind, bool _appendRow)
        : compressionKind(_compressionKind), appendRow(_appendRow) {
    }
};

class UnsafeRowSerializerTest : public ::testing::Test,
                                public pollux::VectorBuilder,
                                public testing::WithParamInterface<TestParam> {
public:
    static std::vector<TestParam> getTestParams() {
        static std::vector<TestParam> testParams = {
            {common::CompressionKind::CompressionKind_NONE, false},
            {common::CompressionKind::CompressionKind_ZLIB, true},
            {common::CompressionKind::CompressionKind_SNAPPY, false},
            {common::CompressionKind::CompressionKind_ZSTD, true},
            {common::CompressionKind::CompressionKind_LZ4, false},
            {common::CompressionKind::CompressionKind_GZIP, true}
        };
        return testParams;
    }

protected:
    static void SetUpTestCase() {
        memory::MemoryManager::testingSetInstance({});
    }

    void SetUp() override {
        pool_ = memory::memoryManager()->addLeafPool();
        deregisterVectorSerde();
        deregisterNamedVectorSerde(VectorSerde::Kind::kCompactRow);
        serializer::spark::UnsafeRowVectorSerde::registerVectorSerde();
        serializer::spark::UnsafeRowVectorSerde::registerNamedVectorSerde();
        ASSERT_EQ(getVectorSerde()->kind(), VectorSerde::Kind::kUnsafeRow);
        ASSERT_EQ(
            getNamedVectorSerde(VectorSerde::Kind::kUnsafeRow)->kind(),
            VectorSerde::Kind::kUnsafeRow);
        appendRow_ = GetParam().appendRow;
        compressionKind_ = GetParam().compressionKind;
        options_ = std::make_unique<VectorSerde::Options>(compressionKind_, 0.8);
    }

    void TearDown() override {
        deregisterVectorSerde();
        deregisterNamedVectorSerde(VectorSerde::Kind::kUnsafeRow);
    }

    void serialize(RowVectorPtr row_vector, std::ostream *output) {
        const auto streamInitialSize = output->tellp();
        const auto numRows = row_vector->size();

        std::vector<IndexRange> ranges(numRows);
        for (int i = 0; i < numRows; i++) {
            ranges[i] = IndexRange{i, 1};
        }

        std::unique_ptr<row::UnsafeRowFast> unsafeRow;
        std::vector<vector_size_t> serializedRowSizes(numRows);
        std::vector<vector_size_t *> serializedRowSizesPtr(numRows);
        std::vector<vector_size_t> rows(numRows);
        std::iota(rows.begin(), rows.end(), 0);
        for (auto i = 0; i < numRows; ++i) {
            serializedRowSizesPtr[i] = &serializedRowSizes[i];
        }
        if (appendRow_) {
            unsafeRow = std::make_unique<row::UnsafeRowFast>(row_vector);
            getVectorSerde()->estimateSerializedSize(
                unsafeRow.get(), rows, serializedRowSizesPtr.data());
        }

        auto arena = std::make_unique<StreamArena>(pool_.get());
        auto rowType = std::dynamic_pointer_cast<const RowType>(row_vector->type());
        auto serializer = getVectorSerde()->createIterativeSerializer(
            rowType, numRows, arena.get(), options_.get());

        if (appendRow_) {
            serializer->append(*unsafeRow, rows, serializedRowSizes);
        } else {
            Scratch scratch;
            serializer->append(
                row_vector, melon::Range(ranges.data(), numRows), scratch);
        }

        auto size = serializer->maxSerializedSize();
        OStreamOutputStream out(output);
        serializer->flush(&out);
        if (!needCompression()) {
            ASSERT_EQ(size, output->tellp() - streamInitialSize);
        } else {
            ASSERT_GT(size, output->tellp() - streamInitialSize);
        }
    }

    std::unique_ptr<ByteInputStream> toByteStream(
        const std::vector<std::string_view> &inputs) {
        std::vector<ByteRange> ranges;
        ranges.reserve(inputs.size());

        for (const auto &input: inputs) {
            ranges.push_back(
                {
                    reinterpret_cast<uint8_t *>(const_cast<char *>(input.data())),
                    (int32_t) input.length(),
                    0
                });
        }
        return std::make_unique<BufferInputStream>(std::move(ranges));
    }

    RowVectorPtr deserialize(
        std::shared_ptr<const RowType> rowType,
        const std::vector<std::string_view> &input) {
        auto byteStream = toByteStream(input);

        RowVectorPtr result;
        getVectorSerde()->deserialize(
            byteStream.get(), pool_.get(), rowType, &result, options_.get());
        return result;
    }

    void testRoundTrip(RowVectorPtr row_vector) {
        std::ostringstream out;
        serialize(row_vector, &out);

        auto rowType = std::dynamic_pointer_cast<const RowType>(row_vector->type());
        auto deserialized = deserialize(rowType, {out.str()});
        test::assertEqualVectors(deserialized, row_vector);
    }

    void
    testSerialize(RowVectorPtr row_vector, int8_t *expectedData, size_t dataSize) {
        std::ostringstream out;
        serialize(row_vector, &out);
        if (!needCompression()) {
            // Check the data after header.
            EXPECT_EQ(
                std::memcmp(expectedData, out.str().data() + kHeaderSize, dataSize),
                0);
        }
    }

    void testDeserialize(
        const std::vector<std::string_view> &input,
        RowVectorPtr expectedVector) {
        if (needCompression()) {
            return;
        }
        // Construct the header to make deserialization work.
        std::vector<std::string_view> uncompressedInput = input;
        char header[kHeaderSize] = {0};
        int32_t uncompressedSize = 0;
        for (const auto &in: input) {
            uncompressedSize += in.size();
        }
        auto *headerPtr = reinterpret_cast<int32_t *>(&header);
        headerPtr[0] = uncompressedSize;
        headerPtr[1] = uncompressedSize;
        header[kHeaderSize - 1] = 0;

        uncompressedInput.insert(
            uncompressedInput.begin(), std::string_view(header, kHeaderSize));
        auto results =
                deserialize(as_row_type(expectedVector->type()), uncompressedInput);
        test::assertEqualVectors(expectedVector, results);
    }

    void
    testDeserialize(int8_t *data, size_t dataSize, RowVectorPtr expectedVector) {
        testDeserialize(
            {std::string_view(reinterpret_cast<const char *>(data), dataSize)},
            expectedVector);
    }

    bool needCompression() {
        return compressionKind_ != common::CompressionKind::CompressionKind_NONE;
    }

    std::shared_ptr<memory::MemoryPool> pool_;

private:
    static constexpr int32_t kHeaderSize = sizeof(int32_t) * 2 + sizeof(char);
    common::CompressionKind compressionKind_;
    std::unique_ptr<VectorSerde::Options> options_;
    bool appendRow_;
};

// These expected binary buffers were samples taken using Spark's java code.
TEST_P(UnsafeRowSerializerTest, tinyint) {
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 123, 0, 0, 0, 0, 0, 0, 0
    };
    auto expected = make_row_vector({make_flat_vector(std::vector<int8_t>{123})});

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);
}

TEST_P(UnsafeRowSerializerTest, bigint) {
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 62, 28, -36, -33, 2, 0, 0, 0
    };
    auto expected =
            make_row_vector({make_flat_vector(std::vector<int64_t>{12345678910})});

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);
}

TEST_P(UnsafeRowSerializerTest, double) {
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 125, 63, 53, 94, -70, 73, -109, 64
    };
    auto expected =
            make_row_vector({make_flat_vector(std::vector<double>{1234.432})});

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);
}

TEST_P(UnsafeRowSerializerTest, boolean) {
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 1, 0, 0, 0, 0, 0, 0, 0
    };
    auto expected = make_row_vector({make_flat_vector(std::vector<bool>{true})});

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);
}

TEST_P(UnsafeRowSerializerTest, string) {
    int8_t data[28] = {
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 5, 0,
        0, 0, 16, 0, 0, 0, 72, 69, 76, 76, 79, 0, 0, 0
    };
    auto expected =
            make_row_vector({make_flat_vector(std::vector<StringView>{"HELLO"})});

    testSerialize(expected, data, 28);
    testDeserialize(data, 28, expected);
}

TEST_P(UnsafeRowSerializerTest, null) {
    int8_t data[20] = {
        0, 0, 0, 16, 1, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0
    };
    auto expected = make_row_vector({
        make_nullable_flat_vector(
            std::vector<std::optional<int64_t> >{std::nullopt})
    });

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);
}

// The data result can be obtained by
// test("decimal serialize") {
//   val d1 = new
//   Decimal().set(BigDecimal("123456789012345678901234.57")).toPrecision(38, 2)
//   val row = InternalRow.apply(d1)
//   val unsafeRow = UnsafeProjection.create(Array[DataType](DecimalType(38,
//   2))).apply(row)
//   assert(unsafeRow.getDecimal(0, 38, 2) === d1)
//   unsafeRow.getBaseObject().asInstanceOf[Array[Byte]].foreach(b => print(b +
//   ", ")) print("\n")
// }
TEST_P(UnsafeRowSerializerTest, decimal) {
    // short decimal
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 62, 28, -36, -33, 2, 0, 0, 0
    };
    auto expected =
            make_row_vector({make_constant<int64_t>(12345678910, 1, DECIMAL(12, 2))});

    testSerialize(expected, data, 20);
    testDeserialize(data, 20, expected);

    // long decimal
    int8_t longData[36] = {
        0, 0, 0, 32, 0, 0, 0, 0, 0, 0, 0, 0,
        11, 0, 0, 0, 16, 0, 0, 0, 10, 54, 76, -104,
        34, 126, -86, 106, -36, -70, -63, 0, 0, 0, 0, 0
    };
    auto longExpected = make_row_vector({
        {
            make_constant<int128_t>(
                HugeInt::build(
                    669260, 10962463713375599297U), // 12345678901234567890123457
                1,
                DECIMAL(38, 2))
        }
    });

    testSerialize(longExpected, longData, 36);
    testDeserialize(longData, 36, longExpected);
}

TEST_P(UnsafeRowSerializerTest, manyRows) {
    int8_t data[140] = {
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0,
        0, 0, 16, 0, 0, 0, 109, 97, 110, 121, 0, 0, 0, 0,
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0,
        0, 0, 16, 0, 0, 0, 114, 111, 119, 115, 0, 0, 0, 0,
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0,
        0, 0, 16, 0, 0, 0, 105, 110, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
        0, 0, 16, 0, 0, 0, 97, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 24, 0, 0, 0, 0, 0, 0, 0, 0, 7, 0,
        0, 0, 16, 0, 0, 0, 112, 97, 121, 108, 111, 97, 100, 0
    };
    auto expected = make_row_vector({
        make_flat_vector(
            std::vector<StringView>{"many", "rows", "in", "a", "payload"})
    });

    testSerialize(expected, data, 140);
    testDeserialize(data, 140, expected);
}

TEST_P(UnsafeRowSerializerTest, splitRow) {
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 62, 28, -36, -33, 2, 0, 0, 0
    };
    auto expected =
            make_row_vector({make_flat_vector(std::vector<int64_t>{12345678910})});

    std::vector<std::string_view> buffers;
    const char *rawData = reinterpret_cast<const char *>(data);

    // Split input row into two buffers.
    buffers = {{rawData, 10}, {rawData + 10, 10}};
    testDeserialize(buffers, expected);

    // Split input row into many buffers.
    buffers = {
        {rawData, 4},
        {rawData + 4, 4},
        {rawData + 8, 4},
        {rawData + 12, 8},
    };
    testDeserialize(buffers, expected);

    // One byte at a time.
    buffers.clear();
    for (size_t i = 0; i < 20; i++) {
        buffers.push_back({rawData + i, 1});
    }
    testDeserialize(buffers, expected);
}

TEST_P(UnsafeRowSerializerTest, incompleteRow) {
    // The test data is for non-compression, and we don't know the compressed size
    // to construct header. If the row is incomplete, readBytes will fail.
    if (needCompression()) {
        return;
    }
    int8_t data[20] = {
        0, 0, 0, 16, 0, 0, 0, 0, 0, 0,
        0, 0, 62, 28, -36, -33, 2, 0, 0, 0
    };
    auto expected =
            make_row_vector({make_flat_vector(std::vector<int64_t>{12345678910})});
    const char *rawData = reinterpret_cast<const char *>(data);

    std::vector<std::string_view> buffers;

    // Cut in the middle of the row.
    buffers = {{rawData, 10}};
    POLLUX_ASSERT_RUNTIME_THROW(
        testDeserialize(buffers, expected), "Unable to read full serialized row");

    // Still incomplete row.
    buffers = {{rawData, 10}, {rawData, 5}};
    POLLUX_ASSERT_RUNTIME_THROW(
        testDeserialize(buffers, expected), "Unable to read full serialized row");

    // Cut right after the row size.
    buffers = {{rawData, 4}};
    POLLUX_ASSERT_RUNTIME_THROW(
        testDeserialize(buffers, expected), "Unable to read full serialized row");

    // Cut in the middle of the `size` integer.
    buffers = {{rawData, 2}};
    POLLUX_ASSERT_RUNTIME_THROW(
        testDeserialize(buffers, expected),
        "(2 vs. 2) Reading past end of BufferInputStream");
}

TEST_P(UnsafeRowSerializerTest, types) {
    auto rowType = ROW(
        {
            BOOLEAN(),
            TINYINT(),
            SMALLINT(),
            INTEGER(),
            BIGINT(),
            REAL(),
            DOUBLE(),
            VARCHAR(),
            TIMESTAMP(),
            DECIMAL(20, 2),
            ROW({VARCHAR(), INTEGER(), DECIMAL(20, 3)}),
            ARRAY(INTEGER()),
            ARRAY(DECIMAL(20, 2)),
            ARRAY(INTEGER()),
            MAP(DECIMAL(20, 3), DECIMAL(20, 3)),
            MAP(VARCHAR(), ARRAY(INTEGER()))
        });

    VectorFuzzer::Options opts;
    opts.vectorSize = 5;
    opts.nullRatio = 0.1;
    opts.dictionaryHasNulls = false;
    opts.stringVariableLength = true;
    opts.stringLength = 20;
    opts.containerVariableLength = false;

    // Spark uses microseconds to store timestamp
    opts.timestampPrecision =
            VectorFuzzer::Options::TimestampPrecision::kMicroSeconds;
    opts.containerLength = 10;

    auto seed = melon::Random::rand32();

    KLOG(ERROR) << "Seed: " << seed;
    SCOPED_TRACE(fmt::format("seed: {}", seed));
    VectorFuzzer fuzzer(opts, pool_.get(), seed);

    auto data = fuzzer.fuzzInputRow(rowType);
    testRoundTrip(data);
}

TEST_P(UnsafeRowSerializerTest, date) {
    auto row_vector = make_row_vector({
        make_flat_vector<int32_t>({0, 1}, DATE()),
    });

    testRoundTrip(row_vector);
}

TEST_P(UnsafeRowSerializerTest, unknown) {
    // UNKNOWN type.
    auto row_vector = make_row_vector({
        BaseVector::create_null_constant(UNKNOWN(), 10, pool()),
    });

    testRoundTrip(row_vector);

    // ARRAY(UNKNOWN) type.
    row_vector = make_row_vector({
        make_array_vector(
            {0, 3, 10, 15},
            BaseVector::create_null_constant(UNKNOWN(), 30, pool())),
    });

    testRoundTrip(row_vector);

    // MAP(BIGINT, UNKNOWN) type.
    row_vector = make_row_vector({
        make_map_vector(
            {0, 3, 7},
            make_flat_vector<int64_t>({1, 2, 3, 4, 5, 6, 7, 8, 9}),
            BaseVector::create_null_constant(UNKNOWN(), 9, pool())),
    });

    testRoundTrip(row_vector);

    // Mix of INTEGER, UNKNOWN and DOUBLE.
    row_vector = make_row_vector({
        make_nullable_flat_vector<int32_t>({1, std::nullopt, 3, 4, 5}),
        BaseVector::create_null_constant(UNKNOWN(), 5, pool()),
        make_nullable_flat_vector<double>(
            {1.1, 2.2, std::nullopt, 4.4, std::nullopt}),
    });
}

TEST_P(UnsafeRowSerializerTest, decimalVector) {
    auto rowVectorDecimal = make_row_vector({
        make_flat_vector<int128_t>(
            {
                0,
                123,
                DecimalUtil::kLongDecimalMin,
                DecimalUtil::kLongDecimalMax,
                HugeInt::build(
                    669260, 10962463713375599297U), // 12345678901234567890123457
            },
            DECIMAL(20, 2))
    });
    testRoundTrip(rowVectorDecimal);

    auto rowVectorArray = make_row_vector({
        make_array_vector(
            {0},
            make_constant<int128_t>(
                HugeInt::build(
                    669260, 10962463713375599297U), // 12345678901234567890123457
                1,
                DECIMAL(20, 2)))
    });

    testRoundTrip(rowVectorArray);
}

TEST_P(UnsafeRowSerializerTest, multiPage) {
    auto input =
            make_row_vector({make_flat_vector(std::vector<int64_t>{12345678910, 123})});
    std::ostringstream out;
    serialize(input, &out);
    serialize(input, &out);
    auto expected = make_row_vector({
        make_flat_vector(
            std::vector<int64_t>{12345678910, 123, 12345678910, 123})
    });
    auto rowType = std::dynamic_pointer_cast<const RowType>(input->type());
    auto deserialized = deserialize(rowType, {out.str()});
    test::assertEqualVectors(deserialized, expected);
}

POLLUX_INSTANTIATE_TEST_SUITE_P(
    UnsafeRowSerializerTest,
    UnsafeRowSerializerTest,
    testing::ValuesIn(UnsafeRowSerializerTest::getTestParams()));
